import torch
from mmpretrain import get_model
from PIL import Image
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

model = get_model('vit-base-p16_sam-pre_3rdparty_sa1b-1024px', pretrained=True)
model.cuda()
# inputs = torch.rand(1, 3, 1024, 1024)

# 定义转换函数，将PIL图像转换为归一化的PyTorch张量
transform = transforms.Compose([
    transforms.Resize((320, 800)),  # 确保图像尺寸正确
    transforms.ToTensor(),           # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 加载图像
image_path = 'aaa.jpg'
image = Image.open(image_path).convert('RGB')  # 确保图像为RGB模式

# 应用转换
inputs = transform(image)

# 如果你需要将它变成(1, 3, 1024, 1024)的batch形式，可以使用unsqueeze(0)
inputs = inputs.unsqueeze(0)
inputs = inputs.cuda() 

# 模型前向传播
out = model(inputs)

# 提取特征
feats = model.extract_feat(inputs)
    
# 打印输出的类型和内容
print("Output type:", type(out))
print("Output content:")
for i, tensor in enumerate(out):
    print(f"  Output {i}:")
    print(f"    Shape: {tensor.shape}")
    print(f"    Gradient function: {tensor.grad_fn}")

# 提取特征

# 打印特征的类型和内容
print("Features type:", type(feats))
print("Features content:")
for i, feat in enumerate(feats):
    print(f"  Feature {i}:")
    print(f"    Shape: {feat.shape}")
    print(f"    Gradient function: {feat.grad_fn}")
    

# 选择要可视化的特征图
feat_to_plot = feats[0][0].mean(dim=0).detach().to('cpu').numpy()  # 计算所有通道的平均值
out_to_plot = out[0][0].mean(dim=0).detach().to('cpu').numpy()     # 同样选择第一个样本的第一个输出

# 绘制特征图
plt.figure(figsize=(10, 5))

# 第一个子图 - 特征图
plt.subplot(1, 2, 1)
plt.imshow(feat_to_plot, cmap='viridis')
plt.title('Feature Map')
plt.colorbar()

# 第二个子图 - 输出张量
plt.subplot(1, 2, 2)
plt.imshow(out_to_plot, cmap='viridis')
plt.title('Model Output')
plt.colorbar()

# 保存图像
plt.savefig('feature_and_output.png')